package projects.ufam.filters;

import ema.Matrix;

public class KalmanFilter {

    /** number of measurement vector dimensions */
    private int mp;
    
    /** number of state vector dimensions */
    private int dp;
    
    /** number of control vector dimensions */
    private int cp;
    
    /** predicted state (x'(k)): x(k) = A*x(k-1) + B*u(k) */
    private double[][] state_pre;           
    
    /** corrected state (x(k)): x(k)=x'(k)+K(k)*(z(k)-H*x'(k)) */
    private double[][] state_post;
    
    /** state transition matrix (A) */
    private double[][] transition_matrix;
    
    /** control matrix (B) (it is not used if there is no control)*/
    private double[][] control_matrix;
    
    /** measurement matrix (H) */    
    private double[][] measurement_matrix;
    
    /** process noise covariance matrix (Q) */
    private double[][] process_noise_cov;
    
    /** measurement noise covariance matrix (R) */
    private double[][] measurement_noise_cov;
    
    /** priori error estimate covariance matrix (P'(k)): P'(k)=A*P(k-1)*At + Q) */
    private double[][] error_cov_pre;
    
    /** Kalman gain matrix (K(k)): K(k)=P'(k)*Ht*inv(H*P'(k)*Ht+R) */
    private double[][] gain;
    
    /** posteriori error estimate covariance matrix (P(k)): P(k)=(I-K(k)*H)*P'(k) */
    private double[][] error_cov_post;
    
    /** temporary matrices */
    //private double[][] temp1, temp2, temp3, temp4, temp5;

    public KalmanFilter( int dynam_params, int measure_params, int control_params ) throws Exception {
        if( dynam_params <= 0 || measure_params <= 0 ) throw new IllegalArgumentException();
        if( control_params < 0 ) control_params = dynam_params;
        dp = dynam_params;
        mp = measure_params;
        cp = control_params;
        state_pre = Matrix.generate(dp, 1); // init by transition _matrix*state_post
        state_post = Matrix.generate(dp, 1); // init by the first measurement!!!
        transition_matrix = Matrix.identity(dp, dp);
        process_noise_cov = Matrix.identity(dp, dp, 1e-3); // 1e-5 (1 in OpenCV)
        measurement_matrix = Matrix.identity(mp, dp); // 1 (0 in OpenCV)
        measurement_noise_cov = Matrix.identity(mp, mp, 1e-1); // 1e-1 (1 in OpenCV)
        error_cov_pre = Matrix.generate(dp, dp); // initialized in Predict
        error_cov_post = Matrix.identity(dp, dp); // 1 (0 in OpenCV)
        gain = Matrix.generate(dp, mp);
        control_matrix = (cp > 0)? Matrix.generate(dp, cp): null;
        /*temp1 = Matrix.generate(dp, dp);
        temp2 = Matrix.generate(mp, dp);
        temp3 = Matrix.generate(mp, mp);
        temp4 = Matrix.generate(mp, dp);
        temp5 = Matrix.generate(mp, 1);*/        
    }
    

    public KalmanFilter(int dynam_params, int measure_params) throws Exception {
        this(dynam_params, measure_params, 0);
    }

    public double[][] predict( double[][] control ) {
        // (1) Project the state ahead (
        state_pre = Matrix.times(transition_matrix,	state_post); // update the state: x'(k) = A*x(k)
        if ( (control != null) && (cp > 0) ) 
        	state_pre = Matrix.gemm(control_matrix, control, state_pre, 1, 1); // x'(k) = x'(k) + B*u(k)
        // (2) Project the error covariance ahead
        double[][] temp1 = Matrix.times( transition_matrix, error_cov_post); //update error covariance matrices: temp1 = A*P(k)
        error_cov_pre = Matrix.gemm( temp1, Matrix.transpose(transition_matrix), process_noise_cov, 1, 1 ); // P'(k) = temp1*At + Q
        return state_pre;
    }
    
    public double[][] predict() {
        return predict(null);
    }

    public double[][] correct(/*final*/ double[][] measurement) {
        // (1) Compute the Kalman gain
    	double[][] temp2 = Matrix.times( measurement_matrix, error_cov_pre ); // temp2 = H*P'(k)        
    	double[][] temp3 = Matrix.gemm( temp2, Matrix.transpose(measurement_matrix), measurement_noise_cov, 1, 1 ); // temp3 = temp2*Ht + R 
    	double[][] temp4 = Matrix.solve( temp3, temp2 ); // temp4 = inv(temp3)*temp2 = Kt(k)
        gain = Matrix.transpose( temp4 ); // K(k)
        
        // (2) Update estimate with measurement z(k)
        double[][] temp5 = Matrix.gemm(measurement_matrix, state_pre, measurement, -1, 1 ); // temp5 = z(k) - H*x'(k) 
        state_post = Matrix.gemm( gain, temp5, state_pre, 1, 1 ); // x(k) = x'(k) + K(k)*temp5 
        
        // (3) Update the error covariance.
        error_cov_post = Matrix.gemm( gain, temp2, error_cov_pre, -1, 1 ); // P(k) = P'(k) - K(k)*temp2
        return state_post;
    }

    public void setState_pre( double[][] state_pre ) {
        this.state_pre = state_pre;
    }

    public double[][] getState_pre() {
        return state_pre;
    }

    public void setState_post( double[][] state_post ) {
        this.state_post = state_post;
    }

    public double[][] getState_post() {
        return state_post;
    }

    public void setTransition_matrix( double[][] transition_matrix ) {
        this.transition_matrix = transition_matrix;
    }

    public double[][] getTransition_matrix() {
        return transition_matrix;
    }

    public void setControl_matrix( double[][] control_matrix) {
        this.control_matrix = control_matrix;
    }

    public double[][] getControl_matrix() {
        return control_matrix;
    }

    public void setMeasurement_matrix( double[][] measurement_matrix ) {
        this.measurement_matrix = measurement_matrix;
    }

    public double[][] getMeasurement_matrix() {
        return measurement_matrix;
    }

    public void setProcess_noise_cov( double[][] process_noise_cov ) {
        this.process_noise_cov = process_noise_cov;
    }

    public double[][] getProcess_noise_cov() {
        return process_noise_cov;
    }

    public void setMeasurement_noise_cov( double[][] measurement_noise_cov ) {
        this.measurement_noise_cov = measurement_noise_cov;
    }

    public double[][] getMeasurement_noise_cov() {
        return measurement_noise_cov;
    }

    public void setError_cov_pre( double[][] error_cov_pre ) {
        this.error_cov_pre = error_cov_pre;
    }

    public double[][] getError_cov_pre() {
        return error_cov_pre;
    }

    public void setGain( double[][] gain ) {
        this.gain = gain;
    }

    public double[][] getGain() {
        return gain;
    }

    public void setError_cov_post( double[][] error_cov_post ) {
        this.error_cov_post = error_cov_post;
    }

    public double[][] getError_cov_post() {
        return error_cov_post;
    }

}
